# 导入相应的库
from tensorflow import keras
from tensorflow.keras.callbacks import ReduceLROnPlateau, ModelCheckpoint
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
import tensorflow as tf
from PIL import Image
import numpy as np
import itertools
import json
import os


# 设置图片的高和宽
from ResNet50.resnet50 import ResNet50
from ResNet50.uitls import plot_confusion_matrix

im_height = 224
im_width = 224

#设置一次训练样本数和迭代次数
batch_size = 64    #一次训练所选取的样本数
epochs = 50   #训练时，使用所有数据集对模型进行一次完整的训练，称为一次epoch

image_path = "../datasets/"  # 猫狗数据集路径
train_dir = "../datasets/train"  # 训练集路径
validation_dir = "../datasets/test"  # 验证集路径

# 定义训练集图像生成器，并进行图像增强
train_image_generator = ImageDataGenerator(rescale=1. / 255,  # 归一化
                                           rotation_range=40,  # 旋转范围
                                           width_shift_range=0.2,  # 水平平移范围
                                           height_shift_range=0.2,  # 垂直平移范围
                                           shear_range=0.2,  # 剪切变换的程度
                                           zoom_range=0.2,  # 剪切变换的程度
                                           horizontal_flip=True,  # 水平翻转
                                           fill_mode='nearest')



#使用图像生成器从文件夹train_dir中读取样本，对标签进行独热编码
train_data_gen = train_image_generator.flow_from_directory(directory=train_dir,  # 从训练集路径读取图片
                                                           batch_size=batch_size,  # 一次训练所选取的样本数
                                                           shuffle=True,  # 打乱标签
                                                           target_size=(im_height, im_width),  # 图片resize到224x224大小
                                                           class_mode='categorical')  # 独热编码

# 训练集样本数:20000
total_train = train_data_gen.n

# 定义验证集图像生成器，并对图像进行预处理
validation_image_generator = ImageDataGenerator(rescale=1./255) # 归一化
# 使用图像生成器从验证集validation_dir中读取样本
val_data_gen = validation_image_generator.flow_from_directory(directory=validation_dir,  # 从验证集路径读取图片
                                                              batch_size=batch_size,  # 一次训练所选取的样本数
                                                              shuffle=False,  # 不打乱标签
                                                              target_size=(im_height, im_width),  # 图片resize到224x224大小
                                                              class_mode='categorical')  # one-hot编码

# 验证集样本数 :5000
total_val = val_data_gen.n

#搭建ResNer网络
model = ResNet50(input_shape=[224,224,3], classes=2)
#加载预训练模型
model.load_weights('./save_weights/resnet50_weights_tf_dim_ordering_tf_kernels.h5', by_name=True, skip_mismatch=True)
# 编译模型
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),  # 使用adam优化器，学习率为1e-3
              loss=tf.keras.losses.CategoricalCrossentropy(from_logits=False),  # 交叉熵损失函数
              metrics=["accuracy"])  # 评价函数

#学习率大小收敛的速度慢，太大会阻碍收敛，造成损失函数在最小值处的震荡，甚至导致发散
# 回调函数:学习率衰减，当模型的性能不提升时应该要降低学习率
reduce_lr = ReduceLROnPlateau(
    monitor='val_loss',  # 需要监视的值
    factor=0.1,  # 学习率衰减为原来的1/10
    patience=2,  # 当patience个epoch过去而模型性能不提升时，学习率减少的动作会被触发
    mode='auto',  # 当监测值为val_acc时，模式应为max，当监测值为val_loss时，模式应为min，在auto模式下，评价准则由被监测值的名字自动推断
    verbose=1  # 如果为True，则为每次更新输出一条消息，默认值:False
)

# 回调函数:保存最优模型
checkpoint = ModelCheckpoint(
    filepath='./save_weights/myResNet16111.h5',  # 保存模型的路径
    monitor='val_accuracy',  # 需要监视的值
    save_weights_only=False,  # 若设置为True，则只保存模型权重，否则将保存整个模型（包括模型结构，配置信息等）
    save_best_only=True,  # 当设置为True时，监测值有改进时才会保存当前的模型
    mode='max',  # 当监测值为val_acc时，模式应为max，当监测值为val_loss时，模式应为min，在auto模式下，评价准则由被监测值的名字自动推断
    period=1  # CheckPoint之间的间隔的epoch数
)

# 开始训练
history = model.fit(x=train_data_gen,  # 输入训练集
                    steps_per_epoch=total_train // batch_size,  # 一个epoch包含的训练步数
                    epochs=epochs,  # 训练模型迭代次数
                    validation_data=val_data_gen,  # 输入验证集
                    validation_steps=total_val // batch_size,  # 一个epoch包含的训练步数
                    callbacks=[checkpoint, reduce_lr])  # 执行回调函数

# 保存训练好的模型
model.save('./save_weights/ResNet50_tf.h5')

#评价模型
# 记录训练集和验证集的准确率和损失值
history_dict = history.history
train_loss = history_dict["loss"]  # 训练集损失值
train_accuracy = history_dict["accuracy"]  # 训练集准确率
val_loss = history_dict["val_loss"]  # 验证集损失值
val_accuracy = history_dict["val_accuracy"]  # 验证集准确率

#画图。测试集和验证集都有
#损失函数图
plt.figure()
plt.plot(range(epochs), train_loss, label='train_loss')
plt.plot(range(epochs), val_loss, label='val_loss')
plt.legend()
plt.xlabel('epochs')
plt.ylabel('loss')

#准确度
plt.figure()
plt.plot(range(epochs), train_accuracy, label='train_accuracy')
plt.plot(range(epochs), val_accuracy, label='val_accuracy')
plt.legend()
plt.xlabel('epochs')
plt.ylabel('accuracy')
plt.show()

# 猫和狗两种标签，存入到labels中
labels = ['cats', 'dogs']

# 预测验证集数据整体准确率
Y_pred = model.predict_generator(val_data_gen, total_val // batch_size + 1)
# 将预测的结果转化为one hit向量
Y_pred_classes = np.argmax(Y_pred, axis=1)
# 计算混淆矩阵
confusion_mtx = confusion_matrix(y_true=val_data_gen.classes, y_pred=Y_pred_classes)
# 绘制混淆矩阵
plot_confusion_matrix(confusion_mtx, normalize=True, target_names=labels)
